/*
Copyright (C) 2013 Matthew Netsch

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License as published by the Free Software Foundation; either
version 2.1 of the License, or (at your option) any later version.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
*/

#include "database.h"
#define _CRT_SECURE_NO_WARNINGS
#include <Windows.h>
#include <WinCrypt.h>
#include <cstdio>

namespace PassMan
{
    /*******************************************************
     * IDatabaseManager
     *******************************************************/

    /**
    *   @brief 
    *       Creates a new database manager of type
    *   @return 
    *       Newly created database manager
    */
    IDatabaseManager* 
    IDatabaseManager::Create(
        DatabaseManagerType type)   ///< Type of database manager
    {
        switch (type)
        {
            case DATABASE_FILE: return new DatabaseManagerFile(); break;
            default:
                break;
        }

        return NULL;
    }

    /**
    *   @brief 
    *       Destructor for IDatabaseManager
    *   @return 
    *       N/A
    */ 
    IDatabaseManager::~IDatabaseManager()
    {
    }

    /*******************************************************
     * IDatabase
     *******************************************************/

    /**
    *   @brief 
    *       Destructor for IDatabase
    *   @return 
    *       N/A
    */ 
    IDatabase::~IDatabase()
    {
    }

    /*******************************************************
     * Database
     *******************************************************/

    /**
    *   @brief 
    *       Constructor for database
    *   @return 
    *       N/A
    */
    Database::Database(
        const char *name,   ///< Name of database to create
        const char *key)    ///< Key of database
        : name(NULL),
          key(NULL),
          pass_num(0),
          pass_size(16),
          passes(NULL),
          observer_num(0),
          observer_size(16),
          observers(NULL)
    {
        setName(name);
        setKey(key);

        passes = (Pass**)malloc(sizeof(Pass*) * pass_size);
        observers = (IObserver**)malloc(sizeof(IObserver*) * observer_size);
    }

    /**
    *   @brief 
    *       Deconstructor for database
    *   @return 
    *       N/A
    */
    Database::~Database()
    {
        if (name)
        {
            free(name);
        }

        if (key)
        {
            free(key);
        }

        if (passes)
        {
            for (int i = 0; i < getNumPass(); i++)
            {
                delete getPass(i);
            }

            free(passes);
        }

        if (observers)
        {
            free(observers);
        }
    }

    /**
    *   @brief 
    *       Creates a clone of the database
    *   @return
    *       Clone of database
    */
    IDatabase*
    Database::clone()
    {
        Database *db = new Database(name, key);

        if (!db)
        {
            return NULL;
        }

        for (int i = 0; i < getNumPass(); i++)
        {
            db->addPass((Pass*)Retrieve(i));
        }

        return db;
    }

    /**
    *   @brief 
    *       Gets the name of the database
    *   @return 
    *       Name of database
    */
    char*
    Database::getName()
    {
        return name;
    }

    /**
    *   @brief 
    *       Gets the key of the database
    *   @return 
    *       Key of database
    */
    char*
    Database::getKey()
    {
        return key;
    }

    /**
    *   @brief 
    *       Sets the name of the database
    *   @return 
    *       N/A
    */
    void
    Database::setName(
        const char *name)   ///< Name to set
    {
        //Free old name
        if (Database::name)
        {
            free(Database::name);
            Database::name = NULL;
        }

        //Set new name
        if (name)
        {
            Database::name = (char*)malloc(sizeof(char) * (strlen(name) + 1));

            if (Database::name)
            {
                strcpy(Database::name, name);
            }
        }

        notify();
    }

    /**
    *   @brief 
    *       Sets the key of the database
    *   @return 
    *       N/A
    */
    void
    Database::setKey(
        const char *key)    ///< Key to set
    {  
        //Free old key
        if (Database::key)
        {
            free(Database::key);
            Database::key = NULL;
        }

        //Set new key
        if (key)
        {
            Database::key = (char*)malloc(sizeof(char) * (strlen(key) + 1));

            if (Database::key)
            {
                strcpy(Database::key, key);
            }
        }

        notify();
    }

    /**
    *   @brief 
    *       Gets the password at index
    *   @return 
    *       Password object
    *       NULL if does not exist
    */
    Pass*
    Database::getPass(
        int index)  ///< Index of password (0 based)
    {
        if (index >= pass_num)
        {
            return NULL;
        }

        return passes[index];
    }

    /**
    *   @brief 
    *       Gets the number of passwords in the database
    *   @return 
    *       Number of passwords
    */
    int
    Database::getNumPass()
    {
        return pass_num;
    }

    /**
    *   @brief 
    *       Adds the password to the database
    *   @return 
    *       true if success
    *   @note
    *       Retains ownership of memory of pass
    *       Fails if pass with name already exists
    */
    bool
    Database::addPass( 
        Pass *pass) ///< Password to add
    {
        if (!pass || !passes)
        {
            return false;
        }

        //Make sure pass name is unique
        if (getPass(pass->getName()))
        {
            return false;
        }

        //Determine if we need to expand list
        if (pass_num >= pass_size)
        {
            pass_size <<= 1;
            passes = (Pass**)realloc(passes, pass_size * sizeof(Pass*));

            if (!passes)
            {
                return false;
            }
        }

        //Add to list
        passes[pass_num++] = pass;
        notify();

        return true;
    }

    /**
    *   @brief 
    *       Removes the password from the database
    *   @return 
    *       true if success
    */
    bool
    Database::removePass( 
        Pass *pass) ///< Password to add
    {
        if (!pass || !passes)
        {
            return false;
        }

        char *pass_name = pass->getName();

        if (!pass_name)
        {
            return false;
        }

        //Find pass index
        int index = -1;

        for (int i = 0; i < pass_num; i++)
        {
            Pass *p = getPass(i);

            if (!p)
            {
                continue;
            }

            char *pname = p->getName();

            if (!pname)
            {
                continue;
            }

            if (strcmp(pname, pass_name) == 0)
            {
                //Pass found, delete and remove from list
                index = i;
                delete p;
                passes[i] = NULL;
            }
        }

        //Fail if doesn't exist
        if (index == -1)
        {
            return false;
        }

        //Push rest of list to the left
        for (int i = index; i < pass_num; i++)
        {
            passes[i] = passes[i + 1];
        }

        passes[--pass_num] = NULL;
        notify();
       
        return true;
    }

    /**
    *   @brief 
    *       Gets the password matching name
    *   @return 
    *       Password with name
    */
    Pass*
    Database::getPass(
        const char *name)   ///< Name of password to get
    {
        if (!name)
        {
            return NULL;
        }

        for (int i = 0; i < getNumPass(); i++)
        {
            Pass *pass = getPass(i);

            if (!pass)
            {
                continue;
            }

            char *pname = pass->getName();

            if (!pname)
            {
                continue;
            }

            if (strcmp(pname, name) == 0)
            {
                return pass;
            }
        }

        return NULL;
    }

    /**
    *   @brief 
    *       Creates a new password in the database
    *   @return 
    *       New password
    */
    IPass* 
    Database::Create(
        const char *name)   ///< Name of password
    {
        if (!name)
        {
            return NULL;
        }

        //Try to allocate pass
        Pass *pass = new Pass(name);

        if (!pass)
        {
            return NULL;
        }

        //Add to list
        if (!addPass(pass))
        {
            delete pass;
            return NULL;
        }

        return new Pass(pass);
    }

    /**
    *   @brief 
    *       Gets password with name from database
    *   @return 
    *       Password representing name
    *       NULL if failure
    */
    IPass* 
    Database::Retrieve(
        const char *name)   ///< Name of password to retrieve
    {
        //Find pass and return
        Pass *pass = getPass(name);

        if (!pass)
        {
            return NULL;
        }
       
        return new Pass(pass);
    }

    
    /**
    *   @brief 
    *       Gets password at index from database
    *   @return 
    *       Password representing name
    *       NULL if failure
    */
    IPass* 
    Database::Retrieve(
        int index)   ///< Index of password to retrieve
    {
        //Find pass and return
        Pass *pass = getPass(index);

        if (!pass)
        {
            return NULL;
        }
       
        return new Pass(pass);
    }

    /**
    *   @brief 
    *       Updates the database with password
    *   @return 
    *       N/A
    */
    void   
    Database::Update(
        IPass *pass)    ///< Password to update
    {
        if (!pass)
        {
            return;
        }

        Pass *new_pass = (Pass*)pass;
        Pass *old_pass = getPass(new_pass->getName());

        //Pass must already exist
        if (!old_pass)
        {
            return;
        }

        //Update state
        old_pass->copy(new_pass);
        notify();
    }

    /**
    *   @brief 
    *       Removes pass from database
    *   @return 
    *       N/A
    */
    void
    Database::Destroy(
        IPass *pass)
    {
        if (!pass)
        {
            return;
        }

        Pass *p = (Pass*)pass;
        removePass(p);
    }

    /**
    *   @brief 
    *       Registers an observer for this database
    *   @return 
    *       N/A
    */
    void
    Database::registerObserver(
        IObserver *observer)    ///< Observer to register
    {
        if (!observers || !observer)
        {
            return;
        }

        //Expand list if needed
        if (observer_num >= observer_size)
        {
            observer_size <<= 1;
            observers = (IObserver**)realloc(observers, sizeof(IObserver*) * observer_size);
        }

        observers[observer_num++] = observer;
    }

    /**
    *   @brief 
    *       Notifies observers of a change in the database
    *   @return 
    *       N/A
    */
    void 
    Database::notify()
    {
        if (!observers)
        {
            return;
        }

        for (int i = 0; i < observer_num; i++)
        {
            if (observers[i])
            {
                observers[i]->notify();
            }
        }
    }

    /*******************************************************
     * DatabaseManagerFile
     *******************************************************/

    /**
    *   @brief 
    *       Constructor for file database manager
    *   @return 
    *       N/A
    */
    DatabaseManagerFile::DatabaseManagerFile()
    {
        if (!CryptAcquireContext(&hProv, NULL, MS_ENH_RSA_AES_PROV, PROV_RSA_AES, CRYPT_VERIFYCONTEXT))
        {
            exit((int)GetLastError());
        }
    }

    /**
    *   @brief 
    *       Destructor for file database manager
    *   @return 
    *       N/A
    */
    DatabaseManagerFile::~DatabaseManagerFile()
    {
        CryptReleaseContext(hProv, 0);
    }

    /**
    *   @brief 
    *       Creates a new database with name
    *   @return 
    *       IDatabase* that represents the created database
    *       NULL if database couldn't be created
    *   @note
    *       Will delete previous database if name already
    *       exists.
    *       Specify NULL name to create a database in
    *       memory only
    */
    IDatabase* 
    DatabaseManagerFile::Create(
        const char *name,   ///< Unique name of database (optional)
        const char *key)    ///< Key that unlocks database (optional)
    {
        //Attempt to allocate the database object
        Database *database = new Database(name, key);

        if (!database)
        {
            return NULL;
        }

        //Only attempt to open database if we have a name
        if (name)
        {
            //Attempt to create a new file
            FILE *pFile = fopen(name, "wb");

            if (!pFile)
            {
                delete database;
                return NULL;
            }

            fclose(pFile);
        }

        //Success
        return database;
    }

    /**
    *   @brief 
    *       Retrieves database of name
    *   @return 
    *       IDatabase* object representing name
    */
    IDatabase* 
    DatabaseManagerFile::Retrieve(
        const char *name,   ///< Unique name of database
        const char *key)    ///< Key that unlocks database
    {
        if (!name || !key)
        {
            return NULL;
        }

        //Attempt to generate encryption key from key
        HCRYPTHASH hHash;
        HCRYPTKEY  hKey;

        //Get Hash object
        if (!CryptCreateHash(hProv, CALG_MD5, 0, 0, &hHash))
        {
            return NULL;
        }

        //Hash the key
        if (!CryptHashData(hHash, (const BYTE*)key, strlen(key), 0))
        {
            CryptDestroyHash(hHash);
            return NULL;
        }

        //Create the encryption key
        if (!CryptDeriveKey(hProv, CALG_AES_128, hHash, 0, &hKey))
        {
            CryptDestroyHash(hHash);
            return NULL;
        }

        //Attempt to allocate database object
        Database *database = new Database(name, key);

        if (!database)
        {
            CryptDestroyKey(hKey);
            CryptDestroyHash(hHash);
            return NULL;
        }

        //Attempt to open database for read
        FILE *pFile = fopen(name, "rb");

        if (!pFile)
        {
            CryptDestroyKey(hKey);
            CryptDestroyHash(hHash);
            delete database;
            return NULL;
        }

        //Read contents
        fseek(pFile, 0, SEEK_END);
        long size = ftell(pFile);
        rewind(pFile);

        char *buf = (char*)malloc(sizeof(char) * size);

        if (!buf)
        {
            CryptDestroyKey(hKey);
            CryptDestroyHash(hHash);
            delete(database);
            fclose(pFile);
            return NULL;
        }

        if (fread(buf, sizeof(char), size, pFile) != size)
        {
            CryptDestroyKey(hKey);
            CryptDestroyHash(hHash);
            delete(database);
            free(buf);
            fclose(pFile);
            return NULL;
        }

        fclose(pFile);

        long bytes_read = 0;

        //Get version
        unsigned int version = *((unsigned int*)&buf[bytes_read]);
        bytes_read += sizeof(unsigned int);

        //Check version
        if (version != PASSMAN_VERSION)
        {
            CryptDestroyKey(hKey);
            CryptDestroyHash(hHash);
            free(buf);
            delete database;
            return NULL;
        }

        //Decrypt data
        DWORD buf_size = size - bytes_read;
        if (!CryptDecrypt(hKey, 0, TRUE, 0, (BYTE*)&buf[bytes_read], &buf_size))
        {
            CryptDestroyKey(hKey);
            CryptDestroyHash(hHash);
            free(buf);
            delete database;
            return NULL;
        }
      
        CryptDestroyKey(hKey);
        CryptDestroyHash(hHash);

        //Get key
        char *fkey = &buf[bytes_read];
        bytes_read += strlen(fkey) + 1;

        if (strcmp(key, fkey) != 0)
        {
            free(buf);
            delete database;
            return NULL;
        }

        //Read contents of file
        while (bytes_read < buf_size)
        {
            char *pname = &buf[bytes_read];
            bytes_read += strlen(pname) + 1;

            if (bytes_read >= buf_size)
            {
                break;
            }

            char *pvalue = &buf[bytes_read];
            bytes_read += strlen(pvalue) + 1;

            Pass *pass = new Pass(pname);

            if (!pass)
            {
                free(buf);
                delete database;
                return NULL;
            }

            pass->setValue(pvalue);
            if (!database->addPass(pass))
            {
                free(buf);
                delete database;
                return NULL;
            }
        }

        free(buf);

        //Success
        return database;
    }

    /**
    *   @brief 
    *       Commits changes to database object
    *   @return 
    *       N/A
    */
    void  
    DatabaseManagerFile::Update(
        IDatabase *database)    ///< Database object to update
    {
        //If no database then nothing to do
        if (!database)
        {
            return;
        }

        //Get the name and key of database
        Database    *db   = (Database*)database;
        const char  *name = db->getName();
        const char  *key  = db->getKey();

        if (!name || !key)
        {
            return;
        }

        //Attempt to generate encryption key from key
        HCRYPTHASH hHash;
        HCRYPTKEY  hKey;

        //Get Hash object
        if (!CryptCreateHash(hProv, CALG_MD5, 0, 0, &hHash))
        {
            return;
        }

        //Hash the key
        if (!CryptHashData(hHash, (const BYTE*)key, strlen(key), 0))
        {
            CryptDestroyHash(hHash);
            return;
        }

        //Create the encryption key
        if (!CryptDeriveKey(hProv, CALG_AES_128, hHash, 0, &hKey))
        {
            CryptDestroyHash(hHash);
            return;
        }

        //Open file for write
        FILE *pFile = fopen(name, "w+b");

        if (!pFile)
        {
            CryptDestroyKey(hKey);
            CryptDestroyHash(hHash);
            return;
        }

        //Write version
        unsigned int version = PASSMAN_VERSION;
        fwrite(&version, sizeof(unsigned int), 1, pFile);
        
        //Calculate data size
        DWORD buf_size = (strlen(key) + 1) * sizeof(char);

        //Write contents
        for (int i = 0; i < db->getNumPass(); i++)
        {
            Pass *pass  = db->getPass(i);

            if (!pass)
            {
                continue;
            }

            char *pname  = pass->getName();
            char *pvalue = pass->getValue(); 

            if (!pname || !pvalue)
            {
                continue;
            }

            buf_size += (strlen(pname) + 1) * sizeof(char);
            buf_size += (strlen(pvalue) + 1) * sizeof(char);
        }

        DWORD buf_length = buf_size;

        //Get required buffer size
        if (!CryptEncrypt(hKey, 0, TRUE, 0, NULL, &buf_size, buf_size))
        {
            CryptDestroyKey(hKey);
            CryptDestroyHash(hHash);
            fclose(pFile);
            return;
        }

        //Malloc to req buffer size
        char *buf = (char*)malloc(buf_size);
        int   bytes_written = 0;

        if (!buf)
        {
            CryptDestroyKey(hKey);
            CryptDestroyHash(hHash);
            fclose(pFile);
            return;
        }

        //Fill in data
        memcpy(&buf[bytes_written], key, strlen(key) + 1);
        bytes_written += strlen(key) + 1;

        //Write contents
        for (int i = 0; i < db->getNumPass(); i++)
        {
            Pass *pass  = db->getPass(i);

            if (!pass)
            {
                continue;
            }

            char *pname  = pass->getName();
            char *pvalue = pass->getValue(); 

            if (!pname || !pvalue)
            {
                continue;
            }

            memcpy(&buf[bytes_written], pname, strlen(pname) + 1);
            bytes_written += strlen(pname) + 1;
            memcpy(&buf[bytes_written], pvalue, strlen(pvalue) + 1);
            bytes_written += strlen(pvalue) + 1;
        }

        //Actually encrypt
        if (!CryptEncrypt(hKey, 0, TRUE, 0, (BYTE*)buf, &buf_length, buf_size))
        {
            CryptDestroyKey(hKey);
            CryptDestroyHash(hHash);
            fclose(pFile);
            free(buf);
            return;
        }

        CryptDestroyKey(hKey);
        CryptDestroyHash(hHash);

        //Write data
        fwrite(buf, buf_size, 1, pFile);
        fclose(pFile);
    }

    /**
    *   @brief 
    *       Destroys database that object represents
    *   @return 
    *       N/A
    */
    void  
    DatabaseManagerFile::Destroy(
        IDatabase *database)    ///< Database object to destroy
    {
        //If no database then nothing to do
        if (!database)
        {
            return;
        }

        //Get the name of database
        Database    *db   = (Database*)database;
        const char  *name = db->getName();

        if (!name)
        {
            return;
        }

        //Delete file
        remove(name);
    }
}